package com.yumeisoft.aitables.util;

import cn.hutool.core.io.resource.ResourceUtil;
import cn.hutool.http.HttpRequest;
import com.google.gson.Gson;
import com.yumeisoft.aitables.entity.LlmAnswerPack;
import lombok.Data;

import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

@Data
public class OllamaClient {
    private String model = "qwen2.5-coder:32b";
    private String endpoint = ResourceUtil.readStr("Ollama.txt", Charset.defaultCharset());

    public String postQuestion(String model, String system, String prompt) {
        Map<String, Object> requestBody = new LinkedHashMap<>();
        requestBody.put("model", model);
        requestBody.put("messages", List.of(
                        Map.of("role", "system", "content", system),
                        Map.of("role", "user", "content", prompt)
                )
        );
        requestBody.put("stream", false);
        requestBody.put("temperature", 0);

        String response = HttpRequest.post(endpoint)
                .header("Content-Type", "application/json")
                .body(new Gson().toJson(requestBody))
                .execute().body();
        return ((Map) new Gson().fromJson(response, Map.class).get("message")).get("content").toString();
    }


    public LlmAnswerPack queryLlm(String systemPromptFile, String context, String question, String lang, List<String> contextAppendFormat) {
        LlmAnswerPack answerPack = new LlmAnswerPack();
        System.out.println("[问]\t" + question);
        String systemPrompt = ResourceUtil.readStr(systemPromptFile, StandardCharsets.UTF_8);
        systemPrompt = systemPrompt.replace("${滚动上下文}", context);
        List<String> answerList = new ArrayList<String>();
        while (answerList.isEmpty()||answerList.size()>contextAppendFormat.size()) {
            String llmAnswer = ask(systemPrompt, question);
            answerPack.setLlmAnswer(llmAnswer);
            System.out.println("[答]\n" + llmAnswer);
            answerList = CodeExtractor.getFromAnswer(llmAnswer, lang);
        }
        StringBuilder contextAppend = new StringBuilder();

        for (int i = 0; i < answerList.size(); i++) {
            String answer = answerList.get(i);
            contextAppend.append(contextAppendFormat.get(i).replace("${answer}", answer));
//            System.out.println(answer);
        }
        System.out.println("---------------------------------------------------------------");
        answerPack.setAnswerList(answerList);
        answerPack.setContextAppend(contextAppend.toString());
        return answerPack;
    }

    public String ask(String system, String prompt) {
        String response = postQuestion(model, system, prompt);
        if (response.contains("<think>")) {
            int thinkEnd = response.indexOf("</think>");
            System.err.println(response.substring("<think>".length(), thinkEnd - 1));
            response = response.substring(thinkEnd + "</think>".length());
            response = response.trim();
        }
        return response;
    }
}
